import torch
import torch.nn as nn
import torch.nn.functional as F

from transformers import AutoTokenizer
from base_model.nlp.base_model.first_model_mla.UsherConfig import UsherConfig
import torch.distributed as dist

from typing import Tuple, Optional, Literal

rank = 0
block_size = 128
import math

attn_impl: Literal["naive", "absorb"] = "absorb"


class ParallelEmbedding(nn.Module):
    """
    Embedding layer with parallelism support across distributed processes.

    Args:
        vocab_size (int): Vocabulary size.
        dim (int): Embedding dimension.
    """

    def __init__(self, vocab_size: int, dim: int):
        super().__init__()
        self.vocab_size = vocab_size
        self.dim = dim
        self.part_vocab_size = vocab_size
        self.vocab_start_idx = rank * self.part_vocab_size
        self.vocab_end_idx = self.vocab_start_idx + self.part_vocab_size
        self.weight = nn.Parameter(torch.empty(self.part_vocab_size, self.dim))
        nn.init.xavier_uniform_(self.weight)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass for parallel embedding layer.

        Args:
            x (torch.Tensor): Input tensor containing token indices.

        Returns:
            torch.Tensor: Embedded representations.

        Raises:
            ValueError: If `world_size` is not defined.
        """
        y = F.embedding(x, self.weight)
        return y


def linear(x: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor] = None) -> torch.Tensor:
    """
    对输入数据应用线性变换：y = xA^T + b。
    该函数支持基于量化和张量格式的专用实现。

    参数:
        x (torch.Tensor): 输入张量。
        weight (torch.Tensor): 权重张量。可能是量化的，在某些情况下需要反量化。
        bias (Optional[torch.Tensor]): 要添加的偏置张量。默认为 None。

    返回:
        torch.Tensor: 线性变换的结果，具体计算可能涉及基于输入参数的量化感知运算。

    注意事项:
        - 如果 `weight` 是量化的（例如 `element_size() == 1`），则使用反量化版本进行计算。
        - 如果 `gemm_impl == "bf16"`，则应用反量化和 `bf16` GEMM 操作。
        - 对于其他情况，函数对 `x` 应用量化，并使用 `fp8_gemm` 进行计算。
    """
    return F.linear(x, weight, bias)


class Linear(nn.Module):
    """
    自定义线性层，支持量化权重和可选的偏置。

    参数:
        in_features (int): 输入特征的数量。
        out_features (int): 输出特征的数量。
        bias (bool): 是否包含偏置项。默认为 False。
        dtype (可选): 层的数据类型。默认为 `torch.bfloat16`。
    """

    def __init__(self, in_features: int, out_features: int, bias: bool = False, dtype=None):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.weight = nn.Parameter(torch.empty(out_features, in_features))
        # 添加初始化
        nn.init.xavier_uniform_(self.weight)  # Xavier初始化

        self.register_parameter("scale", None)
        if bias:
            self.bias = nn.Parameter(torch.empty(out_features))
        else:
            self.register_parameter("bias", None)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        自定义线性层的前向传播。

        参数:
            x (torch.Tensor): 输入张量。

        返回:
            torch.Tensor: 线性变换后的张量。
        """
        return linear(x, self.weight, self.bias)


class ColumnParallelLinear(Linear):
    """
    Linear layer with column parallelism, splitting output features across distributed processes.

    Args:
        in_features (int): Number of input features.
        out_features (int): Total number of output features.
        bias (bool): Whether to include a bias term. Defaults to False.
        dtype (optional): Data type for the layer. Defaults to `torch.bfloat16`.
    """

    def __init__(self, in_features: int, out_features: int, bias: bool = False, dtype=None):
        self.part_out_features = out_features
        super().__init__(in_features, self.part_out_features, bias, dtype)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass for column parallel linear layer.

        Args:
            x (torch.Tensor): Input tensor.

        Returns:
            torch.Tensor: Transformed tensor with column-parallel computation.
        """
        y = linear(x, self.weight, self.bias)
        return y


class RowParallelLinear(Linear):
    """
    Linear layer with row parallelism, splitting input features across distributed processes.

    Args:
        in_features (int): Total number of input features.
        out_features (int): Number of output features.
        bias (bool): Whether to include a bias term. Defaults to False.
        dtype (optional): Data type for the layer. Defaults to `torch.bfloat16`.
    """

    def __init__(self, in_features: int, out_features: int, bias: bool = False, dtype=None):
        self.part_in_features = in_features
        super().__init__(self.part_in_features, out_features, bias, dtype)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass for row parallel linear layer.

        Args:
            x (torch.Tensor): Input tensor.

        Returns:
            torch.Tensor: Transformed tensor with row-parallel computation.
        """
        y = linear(x, self.weight)
        if self.bias is not None:
            y += self.bias
        return y


class RMSNorm(nn.Module):
    """
    Root Mean Square Layer Normalization (RMSNorm).

    Args:
        dim (int): Dimension of the input tensor.
        eps (float): Epsilon value for numerical stability. Defaults to 1e-6.
    """

    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.dim = dim
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x: torch.Tensor):
        """
        Forward pass for RMSNorm.

        Args:
            x (torch.Tensor): Input tensor.

        Returns:
            torch.Tensor: Normalized tensor with the same shape as input.
        """
        return F.rms_norm(x, (self.dim,), self.weight, self.eps)


def precompute_freqs_cis(args: UsherConfig) -> torch.Tensor:
    """
    Precomputes frequency-based complex exponential values for rotary positional embeddings.

    Args:
        args (UsherConfig): Model arguments containing positional embedding parameters.

    Returns:
        torch.Tensor: Precomputed complex exponential values for positional embeddings.
    """
    dim = args.qk_rope_head_dim
    seqlen = args.max_seq_len
    beta_fast = args.beta_fast
    beta_slow = args.beta_slow
    base = args.rope_theta
    factor = args.rope_factor

    def find_correction_dim(num_rotations, dim, base, max_seq_len):
        """
        Computes the correction dimension for a given number of rotations in the rotary positional embedding.

        Args:
            num_rotations (float): Number of rotations to compute the correction for.
            dim (int): Dimensionality of the embedding space.
            base (float): Base value for the exponential computation.
            max_seq_len (int): Maximum sequence length.

        Returns:
            float: The correction dimension based on the input parameters.
        """
        return dim * math.log(max_seq_len / (num_rotations * 2 * math.pi)) / (2 * math.log(base))

    def find_correction_range(low_rot, high_rot, dim, base, max_seq_len):
        """
        Computes the range of correction dimensions for rotary positional embeddings.

        Args:
            low_rot (float): Lower bound for the number of rotations.
            high_rot (float): Upper bound for the number of rotations.
            dim (int): Dimensionality of the embedding space.
            base (float): Base value for the exponential computation.
            max_seq_len (int): Maximum sequence length.

        Returns:
            Tuple[int, int]: The range of correction dimensions (low, high), clamped to valid indices.
        """
        low = math.floor(find_correction_dim(low_rot, dim, base, max_seq_len))
        high = math.ceil(find_correction_dim(high_rot, dim, base, max_seq_len))
        return max(low, 0), min(high, dim - 1)

    def linear_ramp_factor(min, max, dim):
        """
        Computes a linear ramp function used to smooth values between a minimum and maximum range.

        Args:
            min (float): Minimum value for the ramp function.
            max (float): Maximum value for the ramp function.
            dim (int): Dimensionality of the ramp tensor.

        Returns:
            torch.Tensor: A tensor of shape (dim,) with values linearly interpolated between 0 and 1,
                clamped to the range [0, 1].
        """
        if min == max:
            max += 0.001
        linear_func = (torch.arange(dim, dtype=torch.float32) - min) / (max - min)
        ramp_func = torch.clamp(linear_func, 0, 1)
        return ramp_func

    freqs = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))
    if seqlen > args.original_seq_len:
        low, high = find_correction_range(beta_fast, beta_slow, dim, base, args.original_seq_len)
        smooth = 1 - linear_ramp_factor(low, high, dim // 2)
        freqs = freqs / factor * (1 - smooth) + freqs * smooth

    t = torch.arange(seqlen)
    freqs = torch.outer(t, freqs)
    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)
    return freqs_cis


def apply_rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
    """
    Applies rotary positional embeddings to the input tensor.

    Args:
        x (torch.Tensor): Input tensor with positional embeddings to be applied.
        freqs_cis (torch.Tensor): Precomputed complex exponential values for positional embeddings.

    Returns:
        torch.Tensor: Tensor with rotary embeddings applied.
    """
    dtype = x.dtype
    x = torch.view_as_complex(x.float().view(*x.shape[:-1], -1, 2))
    freqs_cis = freqs_cis.view(1, x.size(1), 1, x.size(-1))
    y = torch.view_as_real(x * freqs_cis).flatten(3)
    return y.to(dtype)


class MLA(nn.Module):
    """
    多头注意力层 (MLA)。

    Attributes:
        dim (int): 输入特征的维度。
        n_heads (int): 注意力头的数量。
        n_local_heads (int): 分布式系统中本地处理的注意力头数量。
        q_lora_rank (int): 低秩查询投影的秩（用于LoRA优化）。
        kv_lora_rank (int): 低秩键值投影的秩（用于LoRA优化）。
        qk_nope_head_dim (int): 不带位置嵌入的查询-键投影的维度。
        qk_rope_head_dim (int): 带旋转位置嵌入的查询-键投影的维度。
        qk_head_dim (int): 查询-键投影的总维度。
        v_head_dim (int): 值投影的维度。
        softmax_scale (float): softmax计算中的缩放因子。
    """

    def __init__(self, args: UsherConfig):
        super().__init__()
        self.dim = args.dim
        self.n_heads = args.n_heads
        self.n_local_heads = args.n_heads
        self.q_lora_rank = args.q_lora_rank
        self.kv_lora_rank = args.kv_lora_rank
        self.qk_nope_head_dim = args.qk_nope_head_dim
        self.qk_rope_head_dim = args.qk_rope_head_dim
        self.qk_head_dim = args.qk_nope_head_dim + args.qk_rope_head_dim
        self.v_head_dim = args.v_head_dim

        if self.q_lora_rank == 0:
            # 如果q_lora_rank为0，则直接使用ColumnParallelLinear初始化查询投影层wq
            self.wq = ColumnParallelLinear(self.dim, self.n_heads * self.qk_head_dim)
        else:
            # 否则，使用Linear和RMSNorm初始化低秩查询投影层wq_a和wq_b
            self.wq_a = Linear(self.dim, self.q_lora_rank)
            self.q_norm = RMSNorm(self.q_lora_rank)
            self.wq_b = ColumnParallelLinear(self.q_lora_rank, self.n_heads * self.qk_head_dim)

        # 初始化键值投影层wkv_a和wkv_b，并应用归一化层kv_norm
        self.wkv_a = Linear(self.dim, self.kv_lora_rank + self.qk_rope_head_dim)
        self.kv_norm = RMSNorm(self.kv_lora_rank)
        self.wkv_b = ColumnParallelLinear(self.kv_lora_rank, self.n_heads * (self.qk_nope_head_dim + self.v_head_dim))

        # 初始化输出投影层wo
        self.wo = RowParallelLinear(self.n_heads * self.v_head_dim, self.dim)

        # 计算softmax缩放因子
        self.softmax_scale = self.qk_head_dim ** -0.5
        if args.max_seq_len > args.original_seq_len:
            mscale = 0.1 * args.mscale * math.log(args.rope_factor) + 1.0
            self.softmax_scale = self.softmax_scale * mscale * mscale

        # 根据attn_impl参数初始化缓存张量
        if attn_impl == "naive":
            self.register_buffer("k_cache", torch.zeros(args.max_batch_size, args.max_seq_len, self.n_local_heads, self.qk_head_dim), persistent=False)
            self.register_buffer("v_cache", torch.zeros(args.max_batch_size, args.max_seq_len, self.n_local_heads, self.v_head_dim), persistent=False)
        else:
            self.register_buffer("kv_cache", torch.zeros(args.max_batch_size, args.max_seq_len, self.kv_lora_rank), persistent=False)
            self.register_buffer("pe_cache", torch.zeros(args.max_batch_size, args.max_seq_len, self.qk_rope_head_dim), persistent=False)

    def forward(self, x: torch.Tensor, start_pos: int, freqs_cis: torch.Tensor, mask: Optional[torch.Tensor]):
        """
        多头注意力层的前向传播。

        Attributes:
            x (torch.Tensor): 输入张量，形状为 (batch_size, seq_len, dim)。
            start_pos (int): 缓存中起始位置。
            freqs_cis (torch.Tensor): 预计算的复数指数值，用于旋转位置编码。
            mask (Optional[torch.Tensor]): 掩码张量，用于排除某些位置的注意力。

        Returns:
            torch.Tensor: 输出张量，形状与输入相同。
        """
        bsz, seqlen, _ = x.size()
        end_pos = start_pos + seqlen

        if self.q_lora_rank == 0:
            # 如果q_lora_rank为0，则直接通过wq进行线性变换
            q = self.wq(x)
        else:
            # 否则，通过wq_a、q_norm和wq_b进行低秩查询投影
            q = self.wq_b(self.q_norm(self.wq_a(x)))

        # 将查询张量调整为 (batch_size, seq_len, n_local_heads, qk_head_dim) 形状
        q = q.view(bsz, seqlen, self.n_local_heads, self.qk_head_dim)

        # 将查询张量拆分为非位置感知部分q_nope和旋转位置编码部分q_pe
        q_nope, q_pe = torch.split(q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)

        # 对q_pe应用旋转位置编码
        q_pe = apply_rotary_emb(q_pe, freqs_cis)

        # 计算键值对
        kv = self.wkv_a(x)
        kv, k_pe = torch.split(kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)

        # 对k_pe应用旋转位置编码
        k_pe = apply_rotary_emb(k_pe.unsqueeze(2), freqs_cis)

        if attn_impl == "naive":
            # 如果使用朴素实现方式
            q = torch.cat([q_nope, q_pe], dim=-1)
            kv = self.wkv_b(self.kv_norm(kv))
            kv = kv.view(bsz, seqlen, self.n_local_heads, self.qk_nope_head_dim + self.v_head_dim)
            k_nope, v = torch.split(kv, [self.qk_nope_head_dim, self.v_head_dim], dim=-1)
            k = torch.cat([k_nope, k_pe.expand(-1, -1, self.n_local_heads, -1)], dim=-1)

            # 更新缓存
            # self.k_cache[:bsz, start_pos:end_pos] = k
            # self.v_cache[:bsz, start_pos:end_pos] = v

            # 修改为：
            self.k_cache = self.k_cache.scatter_(dim=1,
                                                 index=torch.arange(start_pos, end_pos).unsqueeze(0).expand(bsz, -1),
                                                 src=k)
            # 修改为：
            self.v_cache = self.v_cache.scatter_(dim=1,
                                                 index=torch.arange(start_pos, end_pos).unsqueeze(0).expand(bsz, -1),
                                                 src=v)

            # 计算注意力分数
            scores = torch.einsum("bshd,bthd->bsht", q, self.k_cache[:bsz, :end_pos]) * self.softmax_scale
        else:
            # 如果使用其他实现方式
            wkv_b = self.wkv_b.weight
            wkv_b = wkv_b.view(self.n_local_heads, -1, self.kv_lora_rank)
            q_nope = torch.einsum("bshd,hdc->bshc", q_nope, wkv_b[:, :self.qk_nope_head_dim])

            # 更新缓存
            with torch.no_grad():
                self.kv_cache[:bsz, start_pos:end_pos] = self.kv_norm(kv)
                self.pe_cache[:bsz, start_pos:end_pos] = k_pe.squeeze(2)

            # 计算注意力分数
            scores = (torch.einsum("bshc,btc->bsht", q_nope, self.kv_cache[:bsz, :end_pos]) +
                      torch.einsum("bshr,btr->bsht", q_pe, self.pe_cache[:bsz, :end_pos])) * self.softmax_scale

        # 应用掩码（如果有）
        if mask is not None:
            scores += mask.unsqueeze(1)

        # 计算softmax
        scores = scores.softmax(dim=-1, dtype=torch.float32).type_as(x)

        if attn_impl == "naive":
            # 如果使用朴素实现方式，计算最终输出
            x = torch.einsum("bsht,bthd->bshd", scores, self.v_cache[:bsz, :end_pos])
        else:
            # 如果使用其他实现方式，计算最终输出
            x = torch.einsum("bsht,btc->bshc", scores, self.kv_cache[:bsz, :end_pos])
            x = torch.einsum("bshc,hdc->bshd", x, wkv_b[:, -self.v_head_dim:])

        # 通过输出投影层返回最终结果
        x = self.wo(x.flatten(2))
        return x


class MLP(nn.Module):
    """
    Multi-Layer Perceptron (MLP) used as a feed-forward layer.

    Attributes:
        w1 (nn.Module): Linear layer for input-to-hidden transformation.
        w2 (nn.Module): Linear layer for hidden-to-output transformation.
        w3 (nn.Module): Additional linear layer for feature transformation.
    """

    def __init__(self, dim: int, inter_dim: int):
        """
        初始化MLP层。

        Args:
            dim (int): 输入和输出张量的维度。
            inter_dim (int): 中间隐藏层的维度。
        """
        super().__init__()
        self.w1 = ColumnParallelLinear(dim, inter_dim)
        self.w2 = RowParallelLinear(inter_dim, dim)
        self.w3 = ColumnParallelLinear(dim, inter_dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass for the MLP layer.

        Args:
            x (torch.Tensor): Input tensor.

        Returns:
            torch.Tensor: Output tensor after MLP computation.
        """
        return self.w2(F.silu(self.w1(x)) * self.w3(x))


class Gate(nn.Module):
    """
    Gate 类是混合专家模型（Mixture-of-Experts, MoE）中的路由机制模块，用于决定输入数据应该被分配到哪些专家（experts）

    Attributes:
        dim (int):        输入特征的维度。
        topk (int):        每个输入激活的顶级专家数量。
        n_groups (int):        路由分组的数量。
        topk_groups (int):        路由中限制的分组数量。
        score_func (str):        评分函数类型，可以是 "softmax" 或 "sigmoid"。
        route_scale (float):        路由权重的缩放因子。
        weight (torch.nn.Parameter):        可学习的权重参数，形状为 (args.n_routed_experts, args.dim)。
        bias (Optional[torch.nn.Parameter]):        可选的偏置项，仅在特定条件下初始化。
    """

    def __init__(self, args: UsherConfig):
        """
        Initializes the Gate module.

        Args:
            args (UsherConfig): Model arguments containing gating parameters.
        """
        super().__init__()
        self.dim = args.dim
        self.topk = args.n_activated_experts
        self.n_groups = args.n_expert_groups
        self.topk_groups = args.n_limited_groups
        self.score_func = args.score_func
        self.route_scale = args.route_scale
        self.weight = nn.Parameter(torch.empty(args.n_routed_experts, args.dim))
        self.bias = nn.Parameter(torch.empty(args.n_routed_experts)) if self.dim == 7168 else None

    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        步骤分解：
            计算评分 (scores)：
                使用 linear 函数计算输入 x 和权重 self.weight 的线性变换结果。
                根据 score_func 的值，选择使用 softmax 或 sigmoid 对评分进行归一化。
            添加偏置：
                如果 bias 存在，则将其加到评分上。
            分组处理（当 n_groups > 1 时）：
                将评分重新调整为 (batch_size, n_groups, ...) 的形状。
                计算每个组的最大评分或前两个评分的和。
                选择评分最高的 topk_groups 组，并屏蔽其他组的评分。
            选择顶级专家：
                在剩余的评分中，选择评分最高的 topk 个专家索引。
            计算权重：
                根据原始评分和选定的索引提取对应的权重。
                如果使用 sigmoid，则对权重进行归一化。
                最后将权重乘以 route_scale 缩放因子。
            返回结果：
                返回路由权重和选定的专家索引。

        Args:
            x (torch.Tensor): Input tensor.

        Returns:
            Tuple[torch.Tensor, torch.Tensor]: Routing weights and selected expert indices.
        """
        scores = linear(x, self.weight)
        if self.score_func == "softmax":
            scores = scores.softmax(dim=-1, dtype=torch.float32)
        else:
            scores = scores.sigmoid()
        original_scores = scores
        if self.bias is not None:
            scores = scores + self.bias
        if self.n_groups > 1:
            scores = scores.view(x.size(0), self.n_groups, -1)
            if self.bias is None:
                group_scores = scores.amax(dim=-1)
            else:
                group_scores = scores.topk(2, dim=-1)[0].sum(dim=-1)
            indices = group_scores.topk(self.topk_groups, dim=-1)[1]
            mask = scores.new_ones(x.size(0), self.n_groups, dtype=bool).scatter_(1, indices, False)
            scores = scores.masked_fill_(mask.unsqueeze(-1), float("-inf")).flatten(1)
        indices = torch.topk(scores, self.topk, dim=-1)[1]
        weights = original_scores.gather(1, indices)
        if self.score_func == "sigmoid":
            weights /= weights.sum(dim=-1, keepdim=True)
        weights *= self.route_scale
        return weights.type_as(x), indices


class Expert(nn.Module):
    """
    Expert 层是混合专家模型（Mixture-of-Experts, MoE）中的一个关键组件，负责执行具体的计算任务。每个 Expert 层可以看作是一个小型的前馈神经网络，用于处理输入数据并生成输出。

    Attributes:
        w1 (Linear 层):
            功能: 将输入张量从 dim 维度映射到 inter_dim 维度。
            实现: 使用自定义的 Linear 层，支持量化权重和可选的偏置项。
        w2 (Linear 层):
            功能: 将中间张量从 inter_dim 维度映射回 dim 维度。
            实现: 同样使用自定义的 Linear 层。
        w3 (Linear 层):
            功能: 提供额外的特征变换，增强模型的表达能力。
            实现: 使用自定义的 Linear 层。

    """

    def __init__(self, dim: int, inter_dim: int):
        """
        Initializes the Expert layer.

        Args:
            dim (int): Input and output dimensionality.
            inter_dim (int): Hidden layer dimensionality.
        """
        super().__init__()
        self.w1 = Linear(dim, inter_dim)
        self.w2 = Linear(inter_dim, dim)
        self.w3 = Linear(dim, inter_dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        步骤分解:
            输入到隐藏层 (w1):
                使用 w1 将输入张量 x 从 dim 维度映射到 inter_dim 维度。
            激活函数 (F.silu):
                应用 SiLU（Sigmoid Linear Unit）激活函数，增强模型的非线性表达能力。
                SiLU 函数定义为 ( \text{SiLU}(x) = x \cdot \sigma(x) )，其中 ( \sigma(x) ) 是 Sigmoid 函数。
            特征变换 (w3):
                使用 w3 对激活后的张量进行额外的线性变换，进一步增强特征表示。
            逐元素乘法:
                将 w2 的输入（即 F.silu(w1(x))）与 w3 的输出进行逐元素乘法操作，结合两种不同的特征变换。
            隐藏层到输出 (w2):
                使用 w2 将最终的中间张量从 inter_dim 维度映射回 dim 维度，生成最终的输出张量。

        Args:
            x (torch.Tensor): Input tensor.

        Returns:
            torch.Tensor: Output tensor after expert computation.
        """
        return self.w2(F.silu(self.w1(x)) * self.w3(x))


class MoE(nn.Module):
    """
    MoE 类是混合专家模型（Mixture-of-Experts, MoE）的核心组件之一，用于在分布式系统中高效地分配和处理输入数据。每个输入数据会被路由到一组选定的专家（Expert 层），这些专家独立地处理输入并生成输出。最后，所有专家的输出被聚合，并与共享专家（shared_experts）的输出相加，生成最终的输出张量。

    Attributes:
        dim (int):
            输入特征的维度。
        n_routed_experts (int):
            模型中总的专家数量。
        n_local_experts (int):
            每个分布式进程本地处理的专家数量。
        n_activated_experts (int):
            每个输入激活的顶级专家数量。
        gate (nn.Module):
            路由机制模块，用于决定输入数据应该被分配到哪些专家。
        experts (nn.ModuleList):
            专家模块的列表，每个专家是一个 Expert 层。
        shared_experts (nn.Module):
            共享的专家模块，应用于所有输入。
    """

    def __init__(self, args: UsherConfig):
        """
        关键点:
            专家数量分配:
                n_routed_experts 确保每个进程处理相同数量的专家。
                n_local_experts 是每个进程本地处理的专家数量。
            专家索引:
                experts_start_idx 和 experts_end_idx 定义了当前进程处理的专家范围。
            路由机制:
                gate 是一个 Gate 模块，用于计算每个输入分配到各个专家的权重和索引。
            专家列表:
                experts 是一个 ModuleList，包含当前进程处理的所有专家。其他进程的专家被设置为 None。
            共享专家:
                shared_experts 是一个 MLP 模块，应用于所有输入，提供共享的特征变换。

        Args:
            args (UsherConfig): Model arguments containing MoE parameters.
        """
        super().__init__()
        self.dim = args.dim
        self.n_routed_experts = args.n_routed_experts
        self.n_local_experts = args.n_routed_experts
        self.n_activated_experts = args.n_activated_experts
        self.experts_start_idx = rank * self.n_local_experts
        self.experts_end_idx = self.experts_start_idx + self.n_local_experts
        self.gate = Gate(args)
        self.experts = nn.ModuleList([Expert(args.dim, args.moe_inter_dim) if self.experts_start_idx <= i < self.experts_end_idx else None
                                      for i in range(self.n_routed_experts)])
        self.shared_experts = MLP(args.dim, args.n_shared_experts * args.moe_inter_dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        步骤分解:
        调整输入形状:
            将输入张量 x 调整为形状 (batch_size * seq_len, dim)，以便处理每个位置的输入。
        计算路由权重和索引:
            使用 gate 模块计算每个输入分配到各个专家的权重 weights 和索引 indices。
        初始化输出张量:
            初始化输出张量 y 为零张量，形状与输入相同。
        统计每个专家的激活次数:
            使用 torch.bincount 统计每个专家被激活的次数 counts。
        处理每个专家:
            遍历当前进程处理的专家范围。
            如果某个专家没有被激活（counts[i] == 0），则跳过。
            否则，获取该专家被激活的输入索引 idx 和对应的权重索引 top。
            使用激活的输入和权重计算专家的输出，并累加到 y 中。
        计算共享专家的输出:
            使用 shared_experts 模块计算共享专家的输出 z。

        合并输出:
            将 y 和 z 相加，并调整形状为原始输入的形状 shape。

        Args:
            x (torch.Tensor): Input tensor.

        Returns:
            torch.Tensor: Output tensor after expert routing and computation.
        """
        shape = x.size()
        x = x.view(-1, self.dim)
        weights, indices = self.gate(x)
        y = torch.zeros_like(x)
        counts = torch.bincount(indices.flatten(), minlength=self.n_routed_experts).tolist()
        for i in range(self.experts_start_idx, self.experts_end_idx):
            if counts[i] == 0:
                continue
            expert = self.experts[i]
            idx, top = torch.where(indices == i)
            y[idx] += expert(x[idx]) * weights[idx, top, None]
        z = self.shared_experts(x)

        return (y + z).view(shape)


class Block(nn.Module):
    """
    Block 类是 Transformer 模型中的一个基本构建块，结合了注意力机制（Attention）和前馈网络（Feed-Forward Network）。

    Attributes:
        attn (nn.Module): 注意力层，具体实现为 MLA（多头注意力层）。
        ffn (nn.Module): 前馈网络层，根据 layer_id 决定是使用 MLP 还是 MoE（混合专家模型）。
        attn_norm (nn.Module): 注意力层的层归一化（Layer Normalization），具体实现为 RMSNorm。
        ffn_norm (nn.Module): 前馈网络层的层归一化，具体实现为 RMSNorm。

    """

    def __init__(self, layer_id: int, args: UsherConfig):
        """
        关键点:
        注意力层 (attn):
            使用 MLA 类实现多头注意力机制。
        前馈网络层 (ffn):
            如果 layer_id 小于 args.n_dense_layers，则使用 MLP 类实现前馈网络。
            否则，使用 MoE 类实现混合专家模型。
        层归一化 (attn_norm 和 ffn_norm):
            使用 RMSNorm 类对注意力层和前馈网络层的输入进行归一化处理。

        Args:
            layer_id (int): Layer index in the transformer.
            args (UsherConfig): Model arguments containing block parameters.
        """
        super().__init__()
        self.attn = MLA(args)
        self.ffn = MLP(args.dim, args.inter_dim) if layer_id < args.n_dense_layers else MoE(args)
        self.attn_norm = RMSNorm(args.dim)
        self.ffn_norm = RMSNorm(args.dim)

    def forward(self, x: torch.Tensor, start_pos: int, freqs_cis: torch.Tensor, mask: Optional[torch.Tensor]) -> torch.Tensor:
        """
        步骤分解:
        注意力层计算:
            首先对输入张量 x 进行层归一化处理 self.attn_norm(x)。
            然后将归一化后的张量传递给注意力层 self.attn，计算注意力机制的结果。
            将注意力层的输出与原始输入 x 相加，实现残差连接（Residual Connection）。
        前馈网络层计算:
            对注意力层的输出进行层归一化处理 self.ffn_norm(x)。
            然后将归一化后的张量传递给前馈网络层 self.ffn，计算前馈网络的结果。
            将前馈网络的输出与注意力层的输出相加，再次实现残差连接。
        返回结果:
            最终返回经过注意力层和前馈网络层处理后的输出张量 x。

        Args:
            x (torch.Tensor): Input tensor.
            start_pos (int): Starting position in the sequence.
            freqs_cis (torch.Tensor): Precomputed complex exponential values for rotary embeddings.
            mask (Optional[torch.Tensor]): Mask tensor to exclude certain positions from attention.

        Returns:
            torch.Tensor: Output tensor after block computation.
        """
        x = x + self.attn(self.attn_norm(x), start_pos, freqs_cis, mask)
        x = x + self.ffn(self.ffn_norm(x))
        return x


class Transformer(nn.Module):
    """
    Transformer 类是整个 Transformer 模型的核心实现，结合了嵌入层、多层 Transformer 块（Block）、归一化层和输出投影层

    Attributes:
        max_seq_len: 模型支持的最大序列长度。
        embed: 并行嵌入层（ParallelEmbedding），用于将输入 token 转换为向量表示。
        layers: 包含多个 Transformer 块（Block）的模块列表，每个块结合了注意力机制和前馈网络。
        norm: 最终的归一化层（RMSNorm），对所有块的输出进行归一化处理。
        head: 输出投影层（ColumnParallelLinear），将模型的隐藏状态映射到词汇表大小的 logits。
        freqs_cis: 预计算的复数指数值，用于旋转位置编码（Rotary Positional Embeddings）。
    """

    def __init__(self, args: UsherConfig):
        """
        关键点:
            分布式设置:
                 rank 用于管理分布式训练环境。
            数据类型:
                根据 args.dtype 设置线性层的数据类型（fp8 或 bfloat16）。
            嵌入层:
                使用 ParallelEmbedding 实现并行化的嵌入层，将输入 token 映射到向量表示。
            Transformer 块:
                构建 args.n_layers 个 Transformer 块（Block），每个块包含注意力机制和前馈网络。
            归一化层:
                使用 RMSNorm 对所有块的输出进行归一化处理。
            输出投影层:
                使用 ColumnParallelLinear 将模型的隐藏状态映射到词汇表大小的 logits。
            旋转位置编码:
                使用 precompute_freqs_cis 预计算旋转位置编码的复数指数值，存储在 freqs_cis 中。

        Args:
            args (UsherConfig): Model arguments containing transformer parameters.
        """
        global rank
        rank = dist.get_rank() if dist.is_initialized() else 0
        super().__init__()
        self.gradient_checkpointing = False
        self.max_seq_len = args.max_seq_len
        self.embed = ParallelEmbedding(args.vocab_size, args.dim)
        self.layers = torch.nn.ModuleList()
        for layer_id in range(args.n_layers):
            self.layers.append(Block(layer_id, args))
        self.norm = RMSNorm(args.dim)
        self.head = ColumnParallelLinear(args.dim, args.vocab_size, dtype=torch.get_default_dtype())
        self.register_buffer("freqs_cis", precompute_freqs_cis(args), persistent=False)

    def enable_gradient_checkpointing(self):
        self.gradient_checkpointing = True
        for layer in self.layers:
            layer.enable_gradient_checkpointing()

    def forward(self, input_ids, start_pos=0, check_labels=None, request_labels=None, **kwargs):
        """
        步骤分解:
            获取序列长度:
                seqlen = tokens.size(1) 获取输入序列的长度。
            嵌入层计算:
                使用 self.embed(tokens) 将输入 token 转换为向量表示。
            旋转位置编码:
                提取与当前序列长度对应的旋转位置编码值 freqs_cis。
            掩码生成:
                如果序列长度大于 1，则生成上三角矩阵作为掩码，防止后续位置的信息泄露。
            Transformer 块计算:
                遍历所有 Transformer 块（self.layers），依次处理输入张量 h。
            归一化层:
                使用 self.norm 对最后一个时间步的隐藏状态进行归一化处理。
            输出投影层:
                使用 self.head 将归一化后的隐藏状态映射到词汇表大小的 logits。

            返回结果:
                返回最终的 logits 张量。

        Args:
            tokens (torch.Tensor): Input tensor of token IDs with shape (batch_size, seq_len).
            start_pos (int, optional): Starting position in the sequence for rotary embeddings. Defaults to 0.

        Returns:
            torch.Tensor: Logits tensor of shape (batch_size, vocab_size).
        """

        seqlen = input_ids.size(1)
        h = self.embed(input_ids)
        freqs_cis = self.freqs_cis[start_pos:start_pos + seqlen]
        mask = None
        if seqlen > 1:
            mask = torch.full((seqlen, seqlen), float("-inf"), device=input_ids.device).triu_(1)

        if self.gradient_checkpointing and self.training:
            from torch.utils.checkpoint import checkpoint
            for i, layer in enumerate(self.layers):
                h = checkpoint(layer, h, start_pos, freqs_cis, mask)
        else:
            for layer in self.layers:
                h = layer(h, start_pos, freqs_cis, mask)

        h = self.norm(h)
        # h = self.norm(h)[:, -1]
        logits = self.head(h)

        return logits

    def generate(self, input_ids, tokenizer: AutoTokenizer, max_length=150, temperature=1.0, top_p=0.9):
        """
        文本生成方法（基于贪心搜索和核采样）

        Args:
            input_ids (torch.Tensor): 输入的token IDs
            max_length (int): 最大生成长度
            temperature (float): 温度参数（控制采样多样性）
            top_p (float): 核采样参数（0.0-1.0）

        Returns:
            torch.Tensor: 生成的token序列
        """
        # 初始化生成序列
        generated = input_ids.clone()
        generated = generated[:-torch.sum(generated == tokenizer.pad_token_id)]
        generated = generated.unsqueeze(0)
        past_key_values = None  # 用于缓存注意力键值对

        for _ in range(max_length):
            # 取最后一个token进行预测
            next_token_logits = self(generated, start_pos=0)  # 调用forward方法
            next_token_logits = next_token_logits[0, -1, :]
            # 应用温度参数
            next_token_logits /= temperature

            # 应用核采样（Top-p sampling）
            filtered_logits = self.top_p_filtering(next_token_logits, top_p=top_p)
            probabilities = F.softmax(filtered_logits, dim=-1)
            next_token = torch.multinomial(probabilities, num_samples=1)

            # 更新生成序列
            generated = torch.cat([generated, next_token.unsqueeze(0)], dim=-1)

            # 检查是否达到终止条件（如生成到EOS_TOKEN）
            if (next_token == tokenizer.eos_token_id).all():
                break

        return generated

    def top_p_filtering(self, logits, top_p=0.9, filter_value=-float('Inf')):
        """
        核采样过滤（Top-p sampling）

        Args:
            logits (torch.Tensor): 模型输出的logits
            top_p (float): 核采样阈值
            filter_value (float): 过滤后的值

        Returns:
            torch.Tensor: 过滤后的logits
        """
        sorted_logits, sorted_indices = torch.sort(logits, descending=True)
        cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)

        # 找到满足累积概率 <= top_p的索引
        sorted_indices_to_remove = cumulative_probs > top_p
        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
        sorted_indices_to_remove[..., 0] = 0

        # 创建过滤掩码
        indices_to_remove = torch.full_like(sorted_indices, False, dtype=bool)
        indices_to_remove.scatter_(
            dim=-1,
            index=sorted_indices,
            src=sorted_indices_to_remove
        )
        logits[indices_to_remove] = filter_value
        return logits


# 示例使用
def modelInit():
    # 设备
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # 配置
    config = UsherConfig()
    # MODEL_NAME = r"D:\lmStudioModel\hugging_model\huggingface\hub\models--deepseek-ai--DeepSeek-R1-Distill-Qwen-7B\snapshots\6602cadec947dbb53e64f3d8d6425320b2197247"

    tokenizer = AutoTokenizer.from_pretrained(config.path)
    tokenizer.save_pretrained(config.path)

    # 模型实例化
    model = Transformer(config).to(device)

    # 保存模型
    torch.save(model.state_dict(), config.path + "/model.pth")
    config.save_pretrained(config.path)
